{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "collapsed": true,
    "pycharm": {
     "is_executing": false
    }
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import tensorflow as tf\n",
    "from tensorflow.keras import layers\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "outputs": [
    {
     "name": "stdout",
     "text": [
      "[[  83   91    1  645 1253  927]\n",
      " [   0   73    8 3215   55  927]\n",
      " [   0    0    0  711  632   71]]\n",
      "[[  83   91    1  645 1253  927]\n",
      " [  73    8 3215   55  927    0]\n",
      " [ 711  632   71    0    0    0]]\n"
     ],
     "output_type": "stream"
    }
   ],
   "source": [
    "raw_inputs = [\n",
    "    [83, 91, 1, 645, 1253, 927],\n",
    "  [73, 8, 3215, 55, 927],\n",
    "  [711, 632, 71]\n",
    "]\n",
    "padded_inptus = tf.keras.preprocessing.sequence.pad_sequences(raw_inputs)\n",
    "print(padded_inptus)\n",
    "padded_inptus = tf.keras.preprocessing.sequence.pad_sequences(raw_inputs,padding='post')\n",
    "print(padded_inptus)"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n",
     "is_executing": false
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "outputs": [
    {
     "name": "stdout",
     "text": [
      "tf.Tensor(\n",
      "[[ True  True  True  True  True  True]\n",
      " [ True  True  True  True  True False]\n",
      " [ True  True  True False False False]], shape=(3, 6), dtype=bool)\n"
     ],
     "output_type": "stream"
    }
   ],
   "source": [
    "embedding = layers.Embedding(input_dim=5000,output_dim=16,mask_zero=True)\n",
    "mask_output = embedding(padded_inptus)\n",
    "print(mask_output._keras_mask)"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n",
     "is_executing": false
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "outputs": [
    {
     "name": "stdout",
     "text": [
      "tf.Tensor(\n",
      "[[ True  True  True  True  True  True]\n",
      " [ True  True  True  True  True False]\n",
      " [ True  True  True False False False]], shape=(3, 6), dtype=bool)\n"
     ],
     "output_type": "stream"
    }
   ],
   "source": [
    "masking_layer = layers.Masking()\n",
    "unmasked_embedding = tf.cast(\n",
    "    tf.tile(tf.expand_dims(padded_inptus,axis=-1),[1,1,16]),\n",
    "    tf.float32\n",
    ")\n",
    "masked_embedding = masking_layer(unmasked_embedding)\n",
    "print(masked_embedding._keras_mask)\n",
    "\n",
    "\n",
    "model = tf.keras.Sequential([\n",
    "    layers.Embedding(input_dim=5000,output_dim=16,mask_zero=True),\n",
    "    layers.LSTM(32)\n",
    "])"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n",
     "is_executing": false
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "outputs": [],
   "source": [
    "inputs = tf.keras.Input(shape=(None,),dtype='int32')\n",
    "x = layers.Embedding(input_dim=5000,output_dim=16,mask_zero=True)(inputs)\n",
    "outputs = layers.LSTM(32)(x)\n",
    "\n",
    "model = tf.keras.Model(inputs,outputs)\n",
    "\n"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n",
     "is_executing": false
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "outputs": [
    {
     "data": {
      "text/plain": "<tf.Tensor: shape=(32, 32), dtype=float32, numpy=\narray([[ 6.5563456e-04, -3.4072709e-05, -2.5979874e-03, ...,\n        -4.6848836e-03, -9.7135443e-04,  2.4124191e-03],\n       [ 3.8403724e-03, -2.1612109e-03,  4.1953321e-03, ...,\n        -2.5533552e-03,  1.2868166e-02,  2.0862322e-03],\n       [-2.2591595e-03,  6.7581227e-03,  1.4892528e-03, ...,\n        -4.1086688e-03,  5.7469070e-04, -2.3707452e-03],\n       ...,\n       [ 7.5938785e-03,  2.6506369e-03,  3.2210480e-03, ...,\n        -1.0068100e-03,  3.5249663e-03, -6.1043091e-03],\n       [ 1.1512681e-03,  5.1162615e-03,  7.4592531e-03, ...,\n        -2.1450960e-03,  1.2224048e-02, -9.0952707e-04],\n       [ 5.0760736e-03, -1.0078367e-02, -3.1843400e-03, ...,\n        -3.7396522e-03, -6.0113623e-05, -3.5711960e-04]], dtype=float32)>"
     },
     "metadata": {},
     "output_type": "execute_result",
     "execution_count": 8
    }
   ],
   "source": [
    "class MyLayer(layers.Layer):\n",
    "    def __init__(self,**kwargs):\n",
    "        super(MyLayer,self).__init__(**kwargs)\n",
    "        self.embedding = layers.Embedding(input_dim=5000,output_dim=16,mask_zero=True)\n",
    "        self.lstm = layers.LSTM(32)\n",
    "    \n",
    "    def call(self,inputs):\n",
    "        x = self.embedding(inputs)\n",
    "        mask = self.embedding.compute_mask(inputs)\n",
    "        output = self.lstm(x,mask=mask)\n",
    "        return output\n",
    "\n",
    "layer = MyLayer()\n",
    "x = np.random.random((32,10)) * 100\n",
    "x = x.astype('int32')\n",
    "layer(x)\n",
    "\n",
    "\n"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n",
     "is_executing": false
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "outputs": [
    {
     "name": "stdout",
     "text": [
      "tf.Tensor(\n",
      "[[ True  True  True]\n",
      " [ True  True  True]\n",
      " [ True  True  True]], shape=(3, 3), dtype=bool)\n",
      "tf.Tensor(\n",
      "[[ True  True  True]\n",
      " [ True  True False]\n",
      " [False False False]], shape=(3, 3), dtype=bool)\n"
     ],
     "output_type": "stream"
    }
   ],
   "source": [
    "class TemporalSplit(tf.keras.layers.Layer):\n",
    "    def call(self,inputs):\n",
    "        return tf.split(inputs,2,axis=1)\n",
    "    \n",
    "    def compute_mask(self,inputs,mask=None):\n",
    "        if mask is None:\n",
    "            return None\n",
    "        return tf.split(mask,2,axis=1)\n",
    "\n",
    "first_half,sencond_half = TemporalSplit()(masked_embedding)\n",
    "print(first_half._keras_mask)\n",
    "print(sencond_half._keras_mask)\n",
    "\n"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n",
     "is_executing": false
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "outputs": [
    {
     "name": "stdout",
     "text": [
      "[[4.37707396 0.26618372 0.684419   2.24155062 3.02215429 2.01368484\n",
      "  7.02290409 4.13425763 2.27202572 7.51708207]\n",
      " [3.46476142 0.21194759 5.62799468 0.2655476  0.48323633 2.80711154\n",
      "  8.48793307 3.77960301 5.94439239 3.93125191]\n",
      " [1.70461702 2.38532908 6.40760015 1.17262807 3.67716665 4.08512562\n",
      "  1.38763645 3.46881381 3.73775593 2.33027969]]\n",
      "tf.Tensor(\n",
      "[[ True False False  True  True  True  True  True  True  True]\n",
      " [ True False  True False False  True  True  True  True  True]\n",
      " [ True  True  True  True  True  True  True  True  True  True]], shape=(3, 10), dtype=bool)\n"
     ],
     "output_type": "stream"
    }
   ],
   "source": [
    "class CustomEmbedding(tf.keras.layers.Layer):\n",
    "    def __init__(self,input_dim,output_dim,mask_zero=False,**kwargs):\n",
    "        super(CustomEmbedding,self).__init__(**kwargs)\n",
    "        self.input_dim = input_dim\n",
    "        self.output_dim = output_dim\n",
    "        self.mask_zero = mask_zero\n",
    "\n",
    "    def build(self,input_shape):\n",
    "        self.embeddings = self.add_weight(\n",
    "            shape=(self.input_dim,self.output_dim),\n",
    "            initializer='random_normal',\n",
    "            dtype='float32'\n",
    "        )\n",
    "    def call(self,inputs):\n",
    "        return tf.nn.embedding_lookup(self.embeddings,inputs)\n",
    "    \n",
    "    def compute_mask(self,inputs,mask=None):\n",
    "        if not self.mask_zero:\n",
    "            return None\n",
    "        return tf.not_equal(inputs,0)\n",
    "\n",
    "layer = CustomEmbedding(10,32,mask_zero=True)\n",
    "x = np.random.random((3,10)) * 9\n",
    "print(x)\n",
    "x = x.astype('int32')\n",
    "\n",
    "y = layer(x)\n",
    "mask = layer.compute_mask(x)\n",
    "print(mask)\n"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n",
     "is_executing": false
    }
   }
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  },
  "pycharm": {
   "stem_cell": {
    "cell_type": "raw",
    "source": [],
    "metadata": {
     "collapsed": false
    }
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}